import argparse
import glob
import cv2
import numpy as np
import os

total_label_1 = 0


def check_path(path):
    # video_path
    if os.path.exists(path):
        pass
    else:
        os.mkdir(path)


def calculate_frame_difference(video_path):
    global total_label_1
    # Load the video
    video = cv2.VideoCapture(video_path)

    # Read the first frame
    _, prev_frame = video.read()

    # Get the video properties
    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))

    # Extract the relevant numbers from the video file name
    file_name = os.path.splitext(os.path.basename(video_path))[0]
    numbers = [int(num) for num in file_name.split('_') if num.isdigit()]
    print(video_path, numbers)
    # Initialize lists to store the frames and labels
    frames = []
    labels = []

    for frame_idx in range(1, total_frames):
        # Read the current frame
        _, curr_frame = video.read()

        # Calculate the difference between the current frame and the previous frame
        frame_diff = cv2.subtract(curr_frame, prev_frame)

        # Convert the difference image to absolute values
        frame_diff = cv2.absdiff(frame_diff, 0)

        # Apply thresholding to create a binary difference image
        _, binary_diff = cv2.threshold(frame_diff, 30, 255, cv2.THRESH_BINARY)

        # Resize the frame
        resized_frame = cv2.resize(binary_diff, (args.length, args.width))

        # Check if the current frame index or the index of the next frame is in the list of numbers
        if frame_idx in numbers or (frame_idx + 1) in numbers or (frame_idx - 1) in numbers:
            label = 1
        else:
            label = 0

        # Append the frame and label to the lists
        frames.append(resized_frame)
        labels.append(label)
        # Update the previous frame for the next iteration
        prev_frame = curr_frame

    indices = [index for index, value in enumerate(labels) if value == 1]
    total_label_1 += len(indices)
    print('total_frames', len(frames), 'label_1 nums', len(indices))

    # Convert the lists to NumPy arrays

    frames = np.array(frames)
    labels = np.array(labels)
    mp4_name = os.path.basename(video_path).replace('.mp4', '')
    frames_path = args.data_path + 'frames/'
    labels_path = args.data_path + 'labels/'
    check_path(frames_path)
    check_path(labels_path)
    frames_np_path = frames_path + mp4_name + "_frames.npy"
    labels_np_path = labels_path + mp4_name + "_labels.npy"
    print(frames_np_path, labels_np_path)
    # Save the frames and labels as NumPy arrays
    np.save(frames_np_path, frames)
    np.save(labels_np_path, labels)

    # Release the video capture
    video.release()


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='This is an example of command-line arguments')

    # Add command-line arguments
    parser.add_argument('--folder_path', type=str,
                        default="forge_video_data")
    parser.add_argument('--data_path', type=str, default='data/')
    parser.add_argument('--length', type=int, default=320)
    parser.add_argument('--width', type=int, default=240)

    # Parse the command-line arguments
    args = parser.parse_args()
    check_path(args.data_path)

    mp4_files = glob.glob(args.folder_path + "/*.mp4")
    #
    for file_path in mp4_files:
        calculate_frame_difference(file_path)
    print('total_label_1:', total_label_1)
